import akshare as ak


def get_a_stock_data(stock_code=None, stock_name=None):
    """
    使用 akshare 获取A股实时行情数据，并支持按股票代码或名称查询。

    Args:
        stock_code (str, optional): 指定的股票代码，如 '600519'。
        stock_name (str, optional): 指定的股票名称，如 '贵州茅台'。

    Returns:
        DataFrame or dict: 如果找到指定股票，则返回该股票的DataFrame或字典数据；
                          如果未指定查询条件，则返回所有A股的DataFrame；
                          如果查询失败或未找到，则返回 None。
    """
    print("正在获取所有A股的实时行情数据...")
    try:
        # 使用 akshare 获取A股所有股票的实时行情，返回一个 DataFrame
        stock_zh_a_spot_df = ak.stock_zh_a_spot_em()

        # 打印数据维度和列名，方便调试
        # print(f"获取到 {len(stock_zh_a_spot_df)} 条股票数据。")
        # print("数据列名:", stock_zh_a_spot_df.columns.tolist())

        # 如果没有指定查询条件，则直接返回整个 DataFrame
        if stock_code is None and stock_name is None:
            return stock_zh_a_spot_df

        # 根据股票代码进行查询
        if stock_code:
            # 使用布尔索引筛选出指定股票代码的数据
            result = stock_zh_a_spot_df[stock_zh_a_spot_df['代码'] == stock_code]
            if not result.empty:
                print(f"成功获取股票代码 {stock_code} 的实时行情。")
                # 返回筛选后的 DataFrame 的第一行数据，并转换为字典，方便查看
                return result.iloc[0].to_dict()
            else:
                print(f"未找到股票代码为 {stock_code} 的数据。")
                return None

        # 根据股票名称进行查询
        if stock_name:
            # 使用布尔索引筛选出指定股票名称的数据
            result = stock_zh_a_spot_df[stock_zh_a_spot_df['名称'] == stock_name]
            if not result.empty:
                print(f"成功获取股票名称 '{stock_name}' 的实时行情。")
                # 返回筛选后的 DataFrame 的第一行数据，并转换为字典
                return result.iloc[0].to_dict()
            else:
                print(f"未找到股票名称为 '{stock_name}' 的数据。")
                return None

    except Exception as e:
        print(f"获取数据时发生错误：{e}")
        return None


# --- 调用示例 ---

# 示例1：查询贵州茅台（股票代码）的实时行情
print("\n--- 示例1：查询股票代码 '600519' 的实时行情 ---")
maotai_data = get_a_stock_data(stock_code='300502')
if maotai_data:
    for key, value in maotai_data.items():
        print(f"{key}: {value}")

# ---
print("\n" + "=" * 50 + "\n")

# 示例2：查询比亚迪（股票名称）的实时行情
# print("--- 示例2：查询股票名称 '比亚迪' 的实时行情 ---")
# byd_data = get_a_stock_data(stock_name='比亚迪')
# if byd_data:
#     for key, value in byd_data.items():
#         print(f"{key}: {value}")
#
# # ---
# print("\n" + "=" * 50 + "\n")
#
# # 示例3：查询一个不存在的股票代码
# print("--- 示例3：查询一个不存在的股票代码 '999999' ---")
# not_exist_data = get_a_stock_data(stock_code='300502')
# # 如果未找到，程序会打印“未找到...”并返回 None
# if not not_exist_data:
#     print("查询结果为 None，符合预期。")